#include<iostream>
#include<unordered_map>
#include<vector>
#include<string>
#include<algorithm>
#include<fstream>
using namespace std;

//树节点
typedef struct HuffmanNode{
    //构造函数 叶子
    HuffmanNode(char k,int w) : key(k), weight(w), left(nullptr), right(nullptr) {};
    //构造函数 非叶
    HuffmanNode(int w) : key('\0'), weight(w), left(nullptr), right(nullptr) {};

    //字符 权重
    char key;
    double weight;
    HuffmanNode* left;
    HuffmanNode* right;
}HuffmanNode;

//节点堆结构 比较函数 小根堆
class ComHuffmanNode{
public:
    bool operator()(const HuffmanNode* n1,const HuffmanNode* n2)
    {
        return n1->weight > n2->weight;
    }
};


class HuffmanTree{
private:
    //存储树节点
    vector<HuffmanNode*> vechufffman;
    //Haffman树根节点
    HuffmanNode* rootNode;
    //编码表
    unordered_map<char,string> haffmanCode;
// 构造函数
public:
    HuffmanTree(){
        rootNode = nullptr;
    }
    ~HuffmanTree(){

    }
//关键函数
public:
    //读入频度文件
    void readData(const string& filename);
    //建造Huffman树
    void buildTree();
    //生成Huffman编码
    void generateCode(HuffmanNode* root, string& code);
    //加密
    string H_encode(const string& str);
    //解密
    string H_decode(const string& str);
//辅助打印函数
public:
    //打印节点数组
    void printVecHuff(const vector<HuffmanNode*>& VH);
    //打印编码
    void printCode(const unordered_map<char,string>& haffmanCode);
};

void HuffmanTree::readData(const string& filename)
{
    fstream file;
    file.open(filename,ios::in);
    if(!file.is_open()){
        cout << "file " << filename << " not exist" << endl;
        exit(-1); 
    }
    //读入叶子 字符+权重
    char key;
    double weight;
    while(file >> key >> weight){
        vechufffman.push_back(new HuffmanNode(key,weight));
    }
    cout << "read from file" << endl;
    this->printVecHuff(this->vechufffman);
}

void HuffmanTree::buildTree()
{
    //建堆
    make_heap(this->vechufffman.begin(),this->vechufffman.end(),ComHuffmanNode());
    //生成堆
    cout << "generate small heap:" << endl;
    this->printVecHuff(this->vechufffman);
    //造树
    while(this->vechufffman.size() > 1){
        //最小值
        HuffmanNode* left = this->vechufffman.front();
        pop_heap(this->vechufffman.begin(),this->vechufffman.end(),ComHuffmanNode());
        this->vechufffman.pop_back();
        //次小值
        HuffmanNode* right = this->vechufffman.front();
        pop_heap(this->vechufffman.begin(),this->vechufffman.end(),ComHuffmanNode());
        this->vechufffman.pop_back();
        //合并成新节点 非叶节点 权值为两者之和
        HuffmanNode* parent = new HuffmanNode(left->weight+right->weight);
        parent->left = left;
        parent->right = right;
        //新节点入堆
        this->vechufffman.push_back(parent);
        push_heap(this->vechufffman.begin(),this->vechufffman.end(),ComHuffmanNode());
    }
    //最后的非叶节点由rootNode继承
    this->rootNode = vechufffman.front();
    //生成编码
    string code = "";
    this->generateCode(this->rootNode,code);
    //打印编码
    cout << "generate the haffman code:" << endl;
    this->printCode(this->haffmanCode);
}

void HuffmanTree::generateCode(HuffmanNode* root, string& code)
{
    //叶子节点
    if(root->left == nullptr){
        this->haffmanCode.insert(pair<char,string>(root->key,code));
        return;
    }
    code.push_back('0');
    generateCode(root->left,code);
    code.pop_back();
    code.push_back('1');
    generateCode(root->right,code);
    code.pop_back();
}

string HuffmanTree::H_encode(const string& str)
{
    string ret = "";
    for(int i=0; i<str.size(); i++){
        ret += this->haffmanCode[str[i]];
    }
    cout << "encode string:" << ret << endl;
    return ret;
}

string HuffmanTree::H_decode(const string& str)
{
    string ret = "";
    HuffmanNode* root = this->rootNode;
    for(int i=0; i<str.length(); i++){
        //寻路
        if(str[i] == '0'){
            root = root->left;
        }else{
            root = root->right;
        }
        //叶节点解密 准备下一次寻路
        if(root->left == nullptr){
            ret.push_back(root->key);
            root = this->rootNode;
        }
    }
    cout << "decode string:" << ret << endl;
    return ret;
}

void HuffmanTree::printVecHuff(const vector<HuffmanNode*>& VH)
{
    for(auto vit = VH.begin(); vit != VH.end(); vit++){
        cout << (*vit)->key << "  " << (*vit)->weight << endl;
    }
}

void HuffmanTree::printCode(const unordered_map<char,string>& haffmanCode)
{
    for(auto mit = haffmanCode.begin(); mit != haffmanCode.end(); mit++){
        cout << mit->first << " : " << mit->second << endl;
    }
}

int main(){
    HuffmanTree* huff = new HuffmanTree();
    huff->readData("frequence.txt");
    huff->buildTree();
    string str = "abcde";
    cout << "to encode:" << str << endl;
    string encode_str = huff->H_encode(str);
    string decode_str = huff->H_decode(encode_str);
    return 0;
}